use axum::{
    extract::ws::{Message, WebSocket, WebSocketUpgrade},
    response::IntoResponse
};

use std::borrow::Cow;
use std::ops::ControlFlow;

use axum::extract::ws::CloseFrame;
use futures_util::{sink::SinkExt, stream::StreamExt};

pub async fn ws_handler( ws: WebSocketUpgrade ) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_socket(socket))
}
async fn handle_socket(mut socket: WebSocket) {
    // 发送一个 ping(某些浏览器不支持) 只是为了启动并获得响应
    if socket.send(Message::Ping(vec![1].into())).await.is_ok() {
        eprintln!("Pinged ...");
    } else {
        eprintln!("Could not send ping!");
        // 这里没有错误，因为我们唯一能做的就是关闭连接。如果不能发送消息，就没有办法挽救状态机
        return;
    }
    // 从客户端接收单个消息(我们可以使用socket接收或发送)。这可能是来自客户端的Ping或hello消息的Pong。
    // 等待来自客户端的消息将阻塞此任务，但不会阻塞其他客户端的任务连接
    if let Some(msg) = socket.recv().await {
        if let Ok(msg) = msg {
            if process_message(msg).is_break() {
                return;
            }
        } else {
            eprintln!("客户端突然断线");
            return;
        }
    }

    // 由于每个客户端都有单独的状态机，我们可以暂停处理在需要时等待一些外部事件(在本例中以sleep表示)。
    // 等待这个客户端完成它的问候语并不会阻止其他客户端连接到服务器并接收他们的问候
    for i in 1..5 {
        if socket.send(Message::Text(format!("Hi {i} times!").into())).await.is_err(){
            eprintln!("客户端突然断线");
            return;
        }
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
    }

    // 通过拆分套接字，我们可以同时发送和接收。在本例中，我们将发送基于某种服务器内部事件(即定时器)向客户端发送未经请求的消息
    let (mut sender, mut receiver) = socket.split();

    // 生成一个任务，将几个消息推送到客户端(不管客户端做什么)
    let mut send_task = tokio::spawn(async move {
        let n_msg = 2;
        for i in 0..n_msg {
            // 如果出现任何websocket错误，我们退出
            if sender.send(Message::Text(format!("Server message {i} ...").into())).await.is_err(){
                return i;
            }

            tokio::time::sleep(std::time::Duration::from_millis(300)).await;
        }

        // 发送关闭消息，并关闭连接 
        if let Err(e) = sender.send(
            Message::Close(Some(CloseFrame { code: axum::extract::ws::close_code::NORMAL,reason: Cow::from("Goodbye") }))
        ).await{
            eprintln!("Could not send Close due to {e}, probably it is ok?");
        }
        n_msg
    });

    // 第二个任务将接收来自客户机的消息并在服务器控制台上打印它们
    let mut recv_task = tokio::spawn(async move {
        let mut cnt = 0;
        while let Some(Ok(msg)) = receiver.next().await {
            cnt += 1;
            // 打印消息，并中断，如果指示这样做
            if process_message(msg).is_break() {
                break;
            }
        }
        cnt
    });

    // 如果其中任何一个任务退出，则中止另一个任务
    tokio::select! {
        rv_a = (&mut send_task) => {
            match rv_a {
                Ok(a) => eprintln!("{a} messages sended"),
                Err(a) => eprintln!("Error sending messages {a:?}")
            }
            recv_task.abort();
        },
        rv_b = (&mut recv_task) => {
            match rv_b {
                Ok(b) => eprintln!("Received {b} messages"),
                Err(b) => eprintln!("Error receiving messages {b:?}")
            }
            send_task.abort();
        }
    }

    // 从处理程序返回将关闭 websocket 连接
    eprintln!("Websocket context destroyed");
}

// 帮助器将消息内容打印到标准输出。对Close有特殊待遇
fn process_message(msg: Message) -> ControlFlow<(), ()> {
    match msg {
        Message::Text(t) => {
            eprintln!("send str: {t:?}");
        }
        Message::Binary(d) => {
            eprintln!("{} send bytes: {:?}", d.len(), d);
        }
        Message::Close(c) => {
            if let Some(cf) = c {
                eprintln!( "send close with code {} and reason `{}`", cf.code, cf.reason);
            } else {
                eprintln!("在没有 CloseFrame 的情况下发送 close 消息");
            }
            return ControlFlow::Break(());
        }

        Message::Pong(v) => {
            eprintln!("sent pong with {v:?}");
        }
        // 你永远不需要手动处理Message::Ping，作为 axum 的 websocket 库将自动为您执行此操作，通过回复 Pong并复制v
        // 但是如果你需要 ping 的内容，你可以在这里看到它们
        Message::Ping(v) => {
            eprintln!("sent ping with {v:?}");
        }
    }
    ControlFlow::Continue(())
}
